# qwen.py

import os
import json
import requests
from datetime import datetime
from ik_analyzer import IKAnalyzer

# ----------------------------
# 配置部分
# ----------------------------
API_KEY = "sk-132e502ffbba4d69b69f5738e308497c"  # 替换为你的实际 API Key
INPUT_DIR = "input"        # 输入文本目录
OUTPUT_DIR = "json"        # 输出 JSON 目录
API_URL = "https://dashscope.aliyuncs.com/compatible-mode/v1/chat/completions"

HEADERS = {
    'Authorization': f'Bearer {API_KEY}',
    'Content-Type': 'application/json'
}

# 创建 IK 分词器实例
ik = IKAnalyzer()

# 创建输出目录（如果不存在）
os.makedirs(OUTPUT_DIR, exist_ok=True)

def call_qwen_api(text):
    """
    调用 Qwen API 解析用户输入文本，提取结构化字段
    """
    data = {
        "model": "qwen-plus",
        "messages": [
            {
                "role": "user",
                "content": text.strip()
            }
        ],
        "tools": [
            {
                "type": "function",
                "function": {
                    "name": "get_fault_detail",
                    "description": "根据故障描述提取其中的信息",
                    "parameters": {
                        "type": "object",
                        "properties": {
                            "productionLine": {"type": "string", "description": "产线"},
                            "batch": {"type": "string", "description": "批次"},
                            "location": {"type": "string", "description": "位置"},
                            "fault": {"type": "string", "description": "故障"}
                        },
                        "required": ["location"]
                    }
                }
            }
        ],
        "tool_choice": "auto"
    }

    try:
        response = requests.post(API_URL, headers=HEADERS, json=data)
        if response.status_code == 200:
            result = response.json()
            tool_response = result.get("choices", [{}])[0].get("message", {}).get("tool_calls", [{}])[0].get("function", {}).get("arguments", {})
            return tool_response
        else:
            print(f"[错误] API 调用失败: 状态码 {response.status_code}, 原因: {response.text}")
            return None
    except Exception as e:
        print(f"[异常] 请求失败: {e}")
        return None

def process_input_files():
    """
    遍历 input/ 下的所有 .txt 文件，调用 Qwen API 并保存结果
    """
    count = 0
    for filename in os.listdir(INPUT_DIR):
        if filename.endswith(".txt"):
            file_path = os.path.join(INPUT_DIR, filename)
            output_filename = f"{os.path.splitext(filename)[0]}.json"
            output_path = os.path.join(OUTPUT_DIR, output_filename)

            with open(file_path, "r", encoding="utf-8") as f:
                text = f.read()

            print(f"[{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] 正在解析文件: {filename}")
            structured_data = call_qwen_api(text)

            if structured_data:
                # 使用 IK 分词器分析原始文本
                smart_tokens = ik.analyze_smart(text)
                max_tokens = ik.analyze_max_word(text)
                keywords = ik.extract_keywords(smart_tokens)

                # 合并数据
                final_data = {
                    "metadata": {
                        "filename": filename,
                        "record_time": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
                    },
                    "structured": structured_data,
                    "ik_analysis": {
                        "smart_tokens": smart_tokens,
                        "max_word_tokens": max_tokens,
                        "keywords": keywords
                    }
                }

                # 写入 JSON 文件
                with open(output_path, "w", encoding="utf-8") as f:
                    json.dump(final_data, f, ensure_ascii=False, indent=4)

                print(f"[保存] 已写入: {output_path}")
                count += 1
            else:
                print(f"[跳过] 无法解析文件: {filename}")

    print(f"\n 完成处理，共成功解析 {count} 个文件。\n")

if __name__ == "__main__":
    process_input_files()